// Copyright 2016 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/task_scheduler/task_scheduler_impl.h"

#include <stddef.h>

#include <string>
#include <utility>
#include <vector>

#include "base/bind.h"
#include "base/bind_helpers.h"
#include "base/callback.h"
#include "base/macros.h"
#include "base/memory/ptr_util.h"
#include "base/synchronization/waitable_event.h"
#include "base/task_scheduler/task_traits.h"
#include "base/task_scheduler/test_task_factory.h"
#include "base/threading/platform_thread.h"
#include "base/threading/simple_thread.h"
#include "base/threading/thread.h"
#include "base/threading/thread_restrictions.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace base {
namespace internal {

    namespace {

        struct TraitsExecutionModePair {
            TraitsExecutionModePair(const TaskTraits& traits,
                ExecutionMode execution_mode)
                : traits(traits)
                , execution_mode(execution_mode)
            {
            }

            TaskTraits traits;
            ExecutionMode execution_mode;
        };

#if ENABLE_THREAD_RESTRICTIONS
        // Returns whether I/O calls are allowed on the current thread.
        bool GetIOAllowed()
        {
            const bool previous_value = ThreadRestrictions::SetIOAllowed(true);
            ThreadRestrictions::SetIOAllowed(previous_value);
            return previous_value;
        }
#endif

        // Verify that the current thread priority and I/O restrictions are appropriate
        // to run a Task with |traits|.
        // Note: ExecutionMode is verified inside TestTaskFactory.
        void VerifyTaskEnvironement(const TaskTraits& traits)
        {
            EXPECT_EQ(traits.priority() == TaskPriority::BACKGROUND
                    ? ThreadPriority::BACKGROUND
                    : ThreadPriority::NORMAL,
                PlatformThread::GetCurrentThreadPriority());

#if ENABLE_THREAD_RESTRICTIONS
            // The #if above is required because GetIOAllowed() always returns true when
            // !ENABLE_THREAD_RESTRICTIONS, even when |traits| don't allow file I/O.
            EXPECT_EQ(traits.with_file_io(), GetIOAllowed());
#endif

            // Verify that the thread the task is running on is named as expected.
            const std::string current_thread_name(PlatformThread::GetName());
            EXPECT_NE(std::string::npos, current_thread_name.find("TaskScheduler"));
            EXPECT_NE(std::string::npos,
                current_thread_name.find(
                    traits.priority() == TaskPriority::BACKGROUND ? "Background"
                                                                  : "Foreground"));
            EXPECT_EQ(traits.with_file_io(),
                current_thread_name.find("FileIO") != std::string::npos);
        }

        void VerifyTaskEnvironementAndSignalEvent(const TaskTraits& traits,
            WaitableEvent* event)
        {
            DCHECK(event);
            VerifyTaskEnvironement(traits);
            event->Signal();
        }

        class ThreadPostingTasks : public SimpleThread {
        public:
            // Creates a thread that posts Tasks to |scheduler| with |traits| and
            // |execution_mode|.
            ThreadPostingTasks(TaskSchedulerImpl* scheduler,
                const TaskTraits& traits,
                ExecutionMode execution_mode)
                : SimpleThread("ThreadPostingTasks")
                , traits_(traits)
                , factory_(scheduler->CreateTaskRunnerWithTraits(traits, execution_mode),
                      execution_mode)
            {
            }

            void WaitForAllTasksToRun() { factory_.WaitForAllTasksToRun(); }

        private:
            void Run() override
            {
                EXPECT_FALSE(factory_.task_runner()->RunsTasksOnCurrentThread());

                const size_t kNumTasksPerThread = 150;
                for (size_t i = 0; i < kNumTasksPerThread; ++i) {
                    factory_.PostTask(test::TestTaskFactory::PostNestedTask::NO,
                        Bind(&VerifyTaskEnvironement, traits_));
                }
            }

            const TaskTraits traits_;
            test::TestTaskFactory factory_;

            DISALLOW_COPY_AND_ASSIGN(ThreadPostingTasks);
        };

        // Returns a vector with a TraitsExecutionModePair for each valid
        // combination of {ExecutionMode, TaskPriority, WithFileIO()}.
        std::vector<TraitsExecutionModePair> GetTraitsExecutionModePairs()
        {
            std::vector<TraitsExecutionModePair> params;

            const ExecutionMode execution_modes[] = { ExecutionMode::PARALLEL,
                ExecutionMode::SEQUENCED,
                ExecutionMode::SINGLE_THREADED };

            for (ExecutionMode execution_mode : execution_modes) {
                for (size_t priority_index = static_cast<size_t>(TaskPriority::LOWEST);
                     priority_index <= static_cast<size_t>(TaskPriority::HIGHEST);
                     ++priority_index) {
                    const TaskPriority priority = static_cast<TaskPriority>(priority_index);
                    params.push_back(TraitsExecutionModePair(
                        TaskTraits().WithPriority(priority), execution_mode));
                    params.push_back(TraitsExecutionModePair(
                        TaskTraits().WithPriority(priority).WithFileIO(), execution_mode));
                }
            }

            return params;
        }

        enum WorkerPoolType {
            BACKGROUND_WORKER_POOL = 0,
            BACKGROUND_FILE_IO_WORKER_POOL,
            FOREGROUND_WORKER_POOL,
            FOREGROUND_FILE_IO_WORKER_POOL,
        };

        size_t GetThreadPoolIndexForTraits(const TaskTraits& traits)
        {
            if (traits.with_file_io()) {
                return traits.priority() == TaskPriority::BACKGROUND
                    ? BACKGROUND_FILE_IO_WORKER_POOL
                    : FOREGROUND_FILE_IO_WORKER_POOL;
            }
            return traits.priority() == TaskPriority::BACKGROUND ? BACKGROUND_WORKER_POOL
                                                                 : FOREGROUND_WORKER_POOL;
        }

        class TaskSchedulerImplTest
            : public testing::TestWithParam<TraitsExecutionModePair> {
        protected:
            TaskSchedulerImplTest() = default;

            void SetUp() override
            {
                using IORestriction = SchedulerWorkerPoolImpl::IORestriction;

                std::vector<TaskSchedulerImpl::WorkerPoolCreationArgs> worker_pools;

                ASSERT_EQ(BACKGROUND_WORKER_POOL, worker_pools.size());
                worker_pools.push_back({ "TaskSchedulerBackground",
                    ThreadPriority::BACKGROUND,
                    IORestriction::DISALLOWED, 1U });

                ASSERT_EQ(BACKGROUND_FILE_IO_WORKER_POOL, worker_pools.size());
                worker_pools.push_back({ "TaskSchedulerBackgroundFileIO",
                    ThreadPriority::BACKGROUND, IORestriction::ALLOWED,
                    3U });

                ASSERT_EQ(FOREGROUND_WORKER_POOL, worker_pools.size());
                worker_pools.push_back({ "TaskSchedulerForeground", ThreadPriority::NORMAL,
                    IORestriction::DISALLOWED, 4U });

                ASSERT_EQ(FOREGROUND_FILE_IO_WORKER_POOL, worker_pools.size());
                worker_pools.push_back({ "TaskSchedulerForegroundFileIO",
                    ThreadPriority::NORMAL, IORestriction::ALLOWED,
                    12U });

                scheduler_ = TaskSchedulerImpl::Create(worker_pools,
                    Bind(&GetThreadPoolIndexForTraits));
                ASSERT_TRUE(scheduler_);
            }

            void TearDown() override { scheduler_->JoinForTesting(); }

            std::unique_ptr<TaskSchedulerImpl> scheduler_;

        private:
            DISALLOW_COPY_AND_ASSIGN(TaskSchedulerImplTest);
        };

    } // namespace

    // Verifies that a Task posted via PostTaskWithTraits with parameterized
    // TaskTraits runs on a thread with the expected priority and I/O restrictions.
    // The ExecutionMode parameter is ignored by this test.
    TEST_P(TaskSchedulerImplTest, PostTaskWithTraits)
    {
        WaitableEvent task_ran(WaitableEvent::ResetPolicy::MANUAL,
            WaitableEvent::InitialState::NOT_SIGNALED);
        scheduler_->PostTaskWithTraits(
            FROM_HERE, GetParam().traits,
            Bind(&VerifyTaskEnvironementAndSignalEvent, GetParam().traits,
                Unretained(&task_ran)));
        task_ran.Wait();
    }

    // Verifies that Tasks posted via a TaskRunner with parameterized TaskTraits and
    // ExecutionMode run on a thread with the expected priority and I/O restrictions
    // and respect the characteristics of their ExecutionMode.
    TEST_P(TaskSchedulerImplTest, PostTasksViaTaskRunner)
    {
        test::TestTaskFactory factory(
            scheduler_->CreateTaskRunnerWithTraits(GetParam().traits,
                GetParam().execution_mode),
            GetParam().execution_mode);
        EXPECT_FALSE(factory.task_runner()->RunsTasksOnCurrentThread());

        const size_t kNumTasksPerTest = 150;
        for (size_t i = 0; i < kNumTasksPerTest; ++i) {
            factory.PostTask(test::TestTaskFactory::PostNestedTask::NO,
                Bind(&VerifyTaskEnvironement, GetParam().traits));
        }

        factory.WaitForAllTasksToRun();
    }

    INSTANTIATE_TEST_CASE_P(OneTraitsExecutionModePair,
        TaskSchedulerImplTest,
        ::testing::ValuesIn(GetTraitsExecutionModePairs()));

    // Spawns threads that simultaneously post Tasks to TaskRunners with various
    // TaskTraits and ExecutionModes. Verifies that each Task runs on a thread with
    // the expected priority and I/O restrictions and respects the characteristics
    // of its ExecutionMode.
    TEST_F(TaskSchedulerImplTest, MultipleTraitsExecutionModePairs)
    {
        std::vector<std::unique_ptr<ThreadPostingTasks>> threads_posting_tasks;
        for (const auto& traits_execution_mode_pair : GetTraitsExecutionModePairs()) {
            threads_posting_tasks.push_back(WrapUnique(new ThreadPostingTasks(
                scheduler_.get(), traits_execution_mode_pair.traits,
                traits_execution_mode_pair.execution_mode)));
            threads_posting_tasks.back()->Start();
        }

        for (const auto& thread : threads_posting_tasks) {
            thread->WaitForAllTasksToRun();
            thread->Join();
        }
    }

    // TODO(fdoray): Add tests with Sequences that move around worker pools once
    // child TaskRunners are supported.

} // namespace internal
} // namespace base
